package SecureUtil;

import java.security.*; 
import java.security.spec.*; 
import javax.crypto.*; 
import javax.crypto.spec.*; 
import javax.crypto.interfaces.*; 
//import com.sun.crypto.provider.SunJCE; 

/** 
* This program executes the Diffie-Hellman key agreement protocol 
* between 2 parties: Client and Server. 
* 
* First call generateFirstHalf() then processFirstHalf then they will generate a shared DES key
* so that we could use that key to encrypt/decrypt the file.
* 
* 
* @author Zhisheng Huang
*/ 

public class DiffieHellman { 
	
	private byte [] _clientPubKeyEnc;
	
	private byte [] _serverPubKeyEnc;
	
	private DHParameterSpec _dhSkipParamSpec; 
	
	private KeyAgreement _clientKeyAgree;
	
	private KeyAgreement _serverKeyAgree;
	
	
	public DiffieHellman() {} 
	
	
	/**
	 * This method return ga mod p, the first step for DH
	 * It's for client's use
	 * 
	 * @return
	 * @throws Exception
	 */
	public byte[] generateFirstHalf() throws Exception {
						
		//generate DH parameters
		AlgorithmParameterGenerator paramGen = AlgorithmParameterGenerator.getInstance("DH"); 
		paramGen.init(512); 
		AlgorithmParameters params = paramGen.generateParameters(); 
		_dhSkipParamSpec = (DHParameterSpec)params.getParameterSpec(DHParameterSpec.class); 
				
		/* 
		* client creates his own DH key pair, using the DH parameters from 
		* above 
		*/ 
		System.out.println("Client: Generate DH keypair ..."); 
		KeyPairGenerator clientKpairGen = KeyPairGenerator.getInstance("DH"); 
		clientKpairGen.initialize(_dhSkipParamSpec); 
		KeyPair clientKpair = clientKpairGen.generateKeyPair(); 
		
		_clientKeyAgree = KeyAgreement.getInstance("DH"); 
		_clientKeyAgree.init(clientKpair.getPrivate()); 
		
		// Client creates and initializes her DH KeyAgreement object 
		System.out.println("Client: Initialize her DH KeyAgreement ..."); 
		KeyAgreement clientKeyAgree = KeyAgreement.getInstance("DH"); 
		clientKeyAgree.init(clientKpair.getPrivate()); 
		
		// Client encodes her public key, and sends it over to Bob. 
		_clientPubKeyEnc = clientKpair.getPublic().getEncoded(); 
		
		return _clientPubKeyEnc;  
	}
	
	
	/**
	 * This method will return the second half of gb mod p in DH algorithm
	 * It's for server use
	 * 
	 * @param receivePubKey
	 * @return
	 * @throws Exception
	 */
	public byte[] processFirstHalf(byte[]receivePubKey ) throws Exception{
		
		KeyFactory serverKeyFac = KeyFactory.getInstance("DH"); 
		X509EncodedKeySpec x509KeySpec = new X509EncodedKeySpec(receivePubKey); 
		PublicKey clientPubKey = serverKeyFac.generatePublic(x509KeySpec); 
		
		/* 
		* server gets the DH parameters associated with client's public key. 
		* He must use the same parameters when he generates his own key 
		* pair. 
		*/ 
		_dhSkipParamSpec = ((DHPublicKey)clientPubKey).getParams(); 
		// server creates his own DH key pair 
		System.out.println("Server: Generate DH keypair ..."); 
		KeyPairGenerator serverKpairGen = KeyPairGenerator.getInstance("DH"); 
		serverKpairGen.initialize(_dhSkipParamSpec); 
		KeyPair serverKpair = serverKpairGen.generateKeyPair(); 
		
		// server creates and initializes his DH KeyAgreement object 
		System.out.println("Server: Initialization ..."); 
		_serverKeyAgree = KeyAgreement.getInstance("DH"); 
		_serverKeyAgree.init(serverKpair.getPrivate()); 
		
		// server encodes his public key, and sends it over to client. 
		_serverPubKeyEnc = serverKpair.getPublic().getEncoded(); 
		
		_serverKeyAgree.doPhase(clientPubKey, true);
		
		return  _serverPubKeyEnc;
	}
	
	/**
	 * 
	 * This method is used in Client to generate a shared DES key 
	 * 
	 * @param serverPubKey
	 * @return
	 * @throws Exception
	 */
	public SecretKey ClientSharedKey(byte [] serverPubKey) throws Exception{
		/* 
		* client uses server's public key for the first (and only) phase 
		* of her version of the DH 
		* protocol. 
		* Before she can do so, she has to instantiate a DH public key 
		* from client's encoded key material. 
		*/ 
		KeyFactory aliceKeyFac = KeyFactory.getInstance("DH"); 
		X509EncodedKeySpec x509KeySpec = new X509EncodedKeySpec(serverPubKey); 
		PublicKey bobPubKey = aliceKeyFac.generatePublic(x509KeySpec); 
		System.out.println("Client: Execute PHASE1 ..."); 
	
		

		
		_clientKeyAgree.doPhase(bobPubKey, true); 
	
		SecretKey clientDesKey = _clientKeyAgree.generateSecret("DES"); 
		
		return clientDesKey;
	
	}
	
	/**
	 * 
	 * This method is used in Server side and it will generate a shared DES key.
	 * 
	 * @return
	 * @throws Exception
	 */
	public SecretKey serverSharedKey() throws Exception{
		
		SecretKey serverDesKey = _serverKeyAgree.generateSecret("DES"); 
		return serverDesKey;
		
	}
	
	public byte[] encryptDH (byte[] plainText, SecretKey sharedKey) throws Exception{
		
		Cipher cipher = Cipher.getInstance("DES/ECB/PKCS5Padding"); 
		cipher.init(Cipher.ENCRYPT_MODE, sharedKey); 
		byte[] ciphertext = cipher.doFinal(plainText); 
		
		return ciphertext;
	}
	
	
	public byte[] decryptDH (byte[]cipherText, SecretKey sharedKey) throws Exception{
		Cipher cipher = Cipher.getInstance("DES/ECB/PKCS5Padding"); 
		cipher.init(Cipher.DECRYPT_MODE, sharedKey); 
		byte[] recovered = cipher.doFinal(cipherText); 
		
		return recovered;
	}
	
	public static void  test(String[]args) throws Exception{
		
		DiffieHellman DHclient = new DiffieHellman();
		DiffieHellman DHserver = new DiffieHellman();
		String plainText = "THis is zhisheng and I love music!!";
		byte[] clientPubKey = DHclient.generateFirstHalf();
		byte[] serverPubKey = DHserver.processFirstHalf(clientPubKey);
		SecretKey clientSharedkey = DHclient.ClientSharedKey(serverPubKey);
		SecretKey serverSharedkey = DHserver.serverSharedKey();
		
		byte [] encryptByte = DHclient.encryptDH(plainText.getBytes(), clientSharedkey);
		
		System.out.println("########  "+new String(plainText.getBytes()));
		
		System.out.println("########  "+ new String(encryptByte));
		
		byte [] recoveryByte = DHserver.decryptDH(encryptByte, serverSharedkey);
		
		System.out.println("########  "+ new String(recoveryByte));
	}

	/** 
	 * Converts a byte to hex digit and writes to the supplied buffer 
	 */ 
	private static void byte2hex(byte b, StringBuffer buf) { 
		char[] hexChars = { '0', '1', '2', '3', '4', '5', '6', '7', '8', 
				'9', 'A', 'B', 'C', 'D', 'E', 'F' }; 
		int high = ((b & 0xf0) >> 4); 
		int low = (b & 0x0f); 
		buf.append(hexChars[high]); 
		buf.append(hexChars[low]); 
	} 
	
	/** 
	 * Converts a byte array to hex string 
	 */ 
	private static String toHexString(byte[] block) { 
		StringBuffer buf = new StringBuffer(); 
	
		int len = block.length; 
	
		for (int i = 0; i < len; i++) { 
			byte2hex(block[i], buf); 
			if (i < len-1) { 
				buf.append(":"); 
			} 
		} 
		return buf.toString(); 
	} 
} 
